module state_mod

  use spark_cubed_sphere
  use params_mod

  implicit none

  private

  public state_type
  public state
  public state_IRK
  public state_IM
  public state_IMEX
  public state_HON
  public state_ref

  type state_panel_array_type
    ! 1: quadrature point
    ! 2: x-axis
    ! 3: y-axis
    ! 4: z-axis
    ! 5: variable
    !                            1 2 3 4 5
    real(r8), pointer, dimension(  :,:,:,:) :: q
    
  end type state_panel_array_type

  type state_type
    logical :: initialized = .false.
    real(r8) :: tm = 0  ! Total mass
    real(r8) :: te = 0  ! Total energy
    type(cubed_sphere_mesh_type), pointer :: mesh => null()
    type(cubed_sphere_array_type) array
    type(state_panel_array_type) domain(6)
  contains
    procedure :: init  => state_init
    procedure :: clear => state_clear
    final :: state_final
  end type state_type

  type(state_type), allocatable :: state(:)
  type(state_type)              :: state_IRK
  type(state_type)              :: state_IM
  type(state_type), allocatable :: state_IMEX(:)
  type(state_type), allocatable :: state_HON(:)
  type(state_type)              :: state_ref

contains

  subroutine state_init(this, mesh)

    class(state_type), intent(inout) :: this
    type(cubed_sphere_mesh_type), intent(in), target :: mesh

    integer ims, ime, jms, jme, kms, kme
    integer igs, ige, jgs, jge

    integer i, idom

    call this%clear()

    this%mesh => mesh
    call this%array%init(mesh)

    if ( model_type == 0 )then
      !!!!!!!!!!!!!!!!!
      ! Add variables !
      !!!!!!!!!!!!!!!!!
      ! ------------------------------------------------------------------------
      !                            Forecast variables
      call this%array%add_var('q1' , &
        dtype='r8', &
        long_name='Metric Jacobian * geopotential depth', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='T', &
        tag='fcst', &
        output='F')

      call this%array%allocate_arrays()
      
      !!!!!!!!!!!!!!!!!
      ! Get variables !
      !!!!!!!!!!!!!!!!!
      do idom = 1, mesh%ndom
        if (mesh%domain(idom)%initialized) then
          call this%array%get_array(this%domain(idom)%q    , idom=idom, loc='CA', tag='fcst')
        end if
      end do
    elseif ( model_type == 1 ) then
      !!!!!!!!!!!!!!!!!
      ! Add variables !
      !!!!!!!!!!!!!!!!!
      ! ------------------------------------------------------------------------
      !                            Forecast variables
      call this%array%add_var('q1' , &
        dtype='r8', &
        long_name='Metric Jacobian * geopotential depth', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='T', &
        tag='fcst', &
        output='F')

      call this%array%add_var('q2' , &
        dtype='r8', &
        long_name='q1 * contravariant wind component along x-axis', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='T', &
        tag='fcst', &
        output='F')

      call this%array%add_var('q3' , &
        dtype='r8', &
        long_name='q1 * contravariant wind component along y-axis', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='T', &
        tag='fcst', &
        output='F')

      call this%array%allocate_arrays()
      
      !!!!!!!!!!!!!!!!!
      ! Get variables !
      !!!!!!!!!!!!!!!!!
      do idom = 1, mesh%ndom
        if (mesh%domain(idom)%initialized) then
          call this%array%get_array(this%domain(idom)%q    , idom=idom, loc='CA', tag='fcst')
        end if
      end do
    elseif(model_type==2)then
      !!!!!!!!!!!!!!!!!
      ! Add variables !
      !!!!!!!!!!!!!!!!!
      ! ------------------------------------------------------------------------
      !                            Forecast variables
      call this%array%add_var('q1' , &
        dtype='r8', &
        long_name='Metric Jacobian * dry density', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='T', &
        tag='fcst', &
        output='F')

      call this%array%add_var('q2' , &
        dtype='r8', &
        long_name='Metric Jacobian * density * contravariant wind component along x-axis', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='T', &
        tag='fcst', &
        output='F')

      call this%array%add_var('q3' , &
        dtype='r8', &
        long_name='Metric Jacobian * density * contravariant wind component along y-axis', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='T', &
        tag='fcst', &
        output='F')
      
      call this%array%add_var('q4' , &
        dtype='r8', &
        long_name='Metric Jacobian * vertical wind', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='T', &
        tag='fcst', &
        output='F')

      call this%array%add_var('q5' , &
        dtype='r8', &
        long_name='q1 * potential temperature', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='T', &
        tag='fcst', &
        output='F')

      call this%array%add_var('q6' , &
        dtype='r8', &
        long_name='q1 * mixing ratio', &
        units='-', &
        loc='CA', &
        with_halo='T', &
        fill_halo='T', &
        tag='fcst', &
        output='F')
      
      call this%array%allocate_arrays()
      
      !!!!!!!!!!!!!!!!!
      ! Get variables !
      !!!!!!!!!!!!!!!!!
      do idom = 1, mesh%ndom
        if (mesh%domain(idom)%initialized) then
          call this%array%get_array(this%domain(idom)%q    , idom=idom, loc='CA', tag='fcst')
        end if
      end do
    end if

    this%initialized = .true.

  end subroutine state_init

  subroutine state_clear(this)

    class(state_type), intent(inout) :: this

    integer idom

    call this%array%clear()

    this%initialized = .false.

  end subroutine state_clear

  subroutine state_final(this)

    type(state_type), intent(inout) :: this

    call this%clear()

  end subroutine state_final

end module state_mod
